#include <iostream>
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <string>
#include <dirent.h>
#include <unistd.h>
#include <vector>
#include <sstream>
#include <fstream>
#include <sys/io.h>
#include <sys/times.h>
#include <iomanip> // setw()
#include <ctime>
using namespace std;

#include "segment_flag.hpp"
using namespace segfg;

void rotateTest();
void detectFromImage();
void detectFromVideo(string strVideoPath);
void generateSample(string strVideoPath);
void for_showing();
void saveImg(int frame);
void legalizationRect(const cv::Mat &src, cv::Rect &rct);

/******************************
*/
int main(int argc, char** argv)
{
	cout << ">> ----" << "\n" << endl;
	if (argc != 3) {
		cout << "error. input parameters" << endl;
		return 0;
	}
	switch (atoi(argv[1])) {
	case 0:{
		detectFromVideo(string(argv[2]));
		break;
	}
	case 1:{
		generateSample(string(argv[2]));
		break;
	}
	}

	cout << "\n" << ">> ----" << endl;
	return 0;
}

// global variables ------------------------------------ //
cv::Mat CurRGB;
cv::Mat CurMono;
cv::Mat ShowImg;
int g_curKeyBoardPushed;
int g_FrmNum = 0;
int g_curFrm = 0;
segfg::SegOCR segocr(segfg::CNN_MODEL_SAPCE2DEPTE);
std::vector<std::vector<cv::Point> > vContourOri;
std::vector<std::vector<cv::Mat> > vOCImg;

void detectFromVideo(string strVideoPath) {
	// openning video
	cv::VideoCapture capture;
	capture = cv::VideoCapture(strVideoPath.c_str());
	double dFrameNum = -1, dFPS = -1, dImgWid = -1, dImgHei = -1;
	if (capture.isOpened()) {
		dFrameNum = capture.get(CV_CAP_PROP_FRAME_COUNT);
		dFPS      = capture.get(CV_CAP_PROP_FPS);
		dImgWid   = capture.get(CV_CAP_PROP_FRAME_WIDTH);
		dImgHei   = capture.get(CV_CAP_PROP_FRAME_HEIGHT);
		g_FrmNum  = (int)dFrameNum;
	} else {
		cout << ">> error. capture.isOpened() return false" << endl;
		return;
	}
	int sta = 80;
	int end = g_FrmNum - 1;
	int step= 5;
	for (;;) {
		if (g_curFrm < sta) g_curFrm = sta;
		if (g_curFrm > end) g_curFrm = end;
		capture.set(CV_CAP_PROP_POS_FRAMES, double(g_curFrm));
		capture.read(CurRGB);
		cvtColor(CurRGB, CurMono, COLOR_BGR2GRAY);
		// process
		segocr.process(CurMono);
		for_showing();
		g_curKeyBoardPushed = waitKey();
		switch (g_curKeyBoardPushed)  {
		case 113: { // q
			g_curFrm -= step;
			if (g_curFrm < sta) g_curFrm = sta;
			break;}
		case 101: { // e
			g_curFrm += step;
			if (g_curFrm > end) g_curFrm = end;
			break;}
		case 27: { 	// esc
			g_curFrm = -1;
			break;}
		case 115: { // s
			saveImg(g_curFrm);
			break;
		}
		}
		if (g_curFrm < 0) break;
	}
	capture.release();
	CurRGB.release();
	CurMono.release();
}

void saveImg(int frame) {
	time_t now = time(0);
	tm *ltm = localtime(&now);
	stringstream ss; 
	ss << 1900 + ltm->tm_year;
	ss << setw(2) << setfill('0') << 1 + ltm->tm_mon;
	ss << setw(2) << setfill('0') << ltm->tm_mday;
	ss << setw(2) << setfill('0') << ltm->tm_hour;
	ss << setw(2) << setfill('0') << ltm->tm_min;
	ss << setw(2) << setfill('0') << ltm->tm_sec;
	ss << setw(5) << setfill('0') << frame;
	int number[3] = {5, 6, 7};
	for (int i = 0; i < (int)vOCImg.size(); i++) {
		for (int j = 0; j < (int)vOCImg[i].size(); j++) {
			stringstream ss1;
			ss1 << setw(4) << setfill('0') << i * j;
			stringstream ss2;
			ss2 << number[j];
			imwrite("./sample/" + ss.str() + ss1.str() + "_" + ss2.str() + ".jpg", vOCImg[i][j]);
		}
	}
	cout << ">> save ok, " << frame << endl;
}

void generateSample(string strVideoPath) {
	// openning video
	cv::VideoCapture capture;
	capture = cv::VideoCapture(strVideoPath.c_str());
	double dFrameNum = -1, dFPS = -1, dImgWid = -1, dImgHei = -1;
	if (capture.isOpened()) {
		dFrameNum = capture.get(CV_CAP_PROP_FRAME_COUNT);
		dFPS      = capture.get(CV_CAP_PROP_FPS);
		dImgWid   = capture.get(CV_CAP_PROP_FRAME_WIDTH);
		dImgHei   = capture.get(CV_CAP_PROP_FRAME_HEIGHT);
		g_FrmNum  = (int)dFrameNum;
	} else {
		cout << ">> error. capture.isOpened() return false" << endl;
		return;
	}
	int sta = 0;
	int end = g_FrmNum - 1;
	int step= 4;
	for (int i = sta; i < end; i += step) {
		capture.set(CV_CAP_PROP_POS_FRAMES, double(i));
		capture.read(CurRGB);
		cvtColor(CurRGB, CurMono, COLOR_BGR2GRAY);
		vContourOri.clear();
		vOCImg.clear();
		segment_flag(CurMono, vContourOri, vOCImg);
		saveImg(i);
		cout << ">> save ok, " << i << endl;
	}
}

void for_showing() 
{
	CurRGB.copyTo(ShowImg);
	stringstream ss;
	ss << g_curFrm;
	ss << " / "; 
	ss << g_FrmNum;
	string str = ss.str();
	cv::putText(ShowImg, str.c_str(), cv::Point(5, 30), 
				FONT_HERSHEY_SIMPLEX, 0.8, 
				Scalar(0, 255, 255, 255), 2, 8);
	for (int i = 0; i < (int)segocr.vContourOri.size(); i++) {
		drawContours(ShowImg, segocr.vContourOri, i, Scalar(255,0,255), 1, LINE_8);
		cv::Mat tmp;
		cv::hconcat(segocr.vOCImg[i], tmp);
		cvtColor(tmp, tmp, COLOR_GRAY2BGR);
		cv::Rect rct = boundingRect(segocr.vContourOri[i]);
		cv::Rect mov(rct.x, rct.y + rct.height + 10, tmp.cols, tmp.rows);
		legalizationRect(ShowImg, mov);
		cv::Mat roi = ShowImg(mov);
		tmp.copyTo(roi);
		
		if ((int)segocr.vCls[i].size() > 0) { // display ocr result
			stringstream ss1;
			for (int j = 0; j < (int)segocr.vCls[i].size(); j++) {
				if (segocr.vCls[i][j] >= 0) {
					ss1 << segocr.vCls[i][j];
				}
				else
					ss1 << "*";
				cout << i << ", " << j << ", " << segocr.vCls[i][j] << endl;
			}
			cv::putText(ShowImg, ss1.str(), cv::Point(mov.x, mov.y + mov.height + 20), 
			FONT_HERSHEY_SIMPLEX, 0.8, Scalar(255, 0, 0), 2, 8);
		}
	}
	
	imshow("win1", ShowImg);
}

void legalizationRect(const cv::Mat &src, cv::Rect &rct) {
	if (rct.x < 0) rct.x = 0;
	if (rct.y < 0) rct.y = 0;
	if (rct.x + rct.width >= src.cols) {
		int sub = abs(rct.x + rct.width - src.cols);
		rct.x -= (sub + 1);
	}
	if (rct.y + rct.height >= src.rows) {
		int sub = abs(rct.y + rct.height - src.rows);
		rct.y -= (sub + 1);
	}
//	cout << src.cols << ", " << src.rows << endl;
//	cout << rct << endl;
}

void detectFromImage() {
	cv::Mat src = imread("./data/q_1.jpg", IMREAD_GRAYSCALE);
	std::vector<std::vector<cv::Point> > vContourOri;
	std::vector<std::vector<cv::Mat> > vOCImg;
	segment_flag(src, vContourOri, vOCImg);
}

/******************************
*/
void rotateTest() {
	cv::Mat src(cv::Size(240, 160), CV_8UC1, Scalar::all(0));
	int edge = 50;
	cv::Rect rct(edge, edge, src.cols - edge * 2, src.rows - edge* 2);
	cv::rectangle(src, rct, cv::Scalar(255), -1, 8, 0);
	cv::rectangle(src, Rect(rct.x + 5, rct.y + 5, rct.width - 10, 4), 
	cv::Scalar(0, 0, 0), -1, 8, 0);
	for (int i = 0; i < 12; i++) {
		cv::Mat rotateImg;
		if (1) {
			Point2f center(src.cols / 2, src.rows / 2);
			float angle = 30.f * i;
			cv::Mat rot = cv::getRotationMatrix2D(center, angle, 1.f);
			cv::RotatedRect box(center, cv::Size(src.rows, src.cols), angle);
			cv::Rect bbox = box.boundingRect();
			rot.at<double>(0, 2) += bbox.width / 2.f - center.x;
			rot.at<double>(1, 2) += bbox.height/ 2.f - center.y;
			cv::warpAffine(src, rotateImg, rot, cv::Size(bbox.width, bbox.height));
		}
		
		vector<vector<Point> > contours;
		vector<Vec4i> hierarchy;
		findContours(rotateImg, contours, hierarchy, RETR_TREE, CHAIN_APPROX_SIMPLE);
		cv::RotatedRect box = minAreaRect(contours[0]);
		cout << setw(5) << (int)box.angle << "    "
		     << setw(5) << (int)box.size.width << "    "
		     << setw(5) << (int)box.size.height << endl;
		
/*		imshow("win", rotateImg);
		waitKey(0);*/
	}
}


